#include <iostream>
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <sys/time.h>
#include "mkl.h"
#include <omp.h>
using namespace std;
double ***A;
double * Alin;

#pragma oss task inout([NB][NB]T)/*为什么这样写也对？*/
//#pragma oss task inout([NB*NB]T)
void dpotrf_tile(double *T, int NB)
{
    int info;
    dpotrf2("L", &NB, T, &NB, &info);
}

#pragma oss task in([NB][NB]T, [NB][NB]B) inout([NB][NB]C)
//#pragma oss task in([NB*NB]T, [NB*NB]B) inout([NB*NB]C)
void dgemm_tile(double *T, double *B, double *C, int NB)
{
    double const done=1.0, dmone=-1.0;
    cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans, NB, NB, NB, dmone, T, NB, B, NB, done, C, NB);
}

#pragma oss task in([NB][NB]T) inout([NB][NB]B)
//#pragma oss task in([NB*NB]T) inout([NB*NB]B)
void dtrsm_tile(double *T, double *B, int NB)
{
    //printf("%d\n",omp_get_num_threads());
    double const done=1.0;
    cblas_dtrsm(CblasColMajor, CblasRight, CblasLower, CblasTrans, CblasNonUnit, NB, NB, done, T, NB, B, NB);
}

#pragma oss task in([NB][NB]T) inout([NB][NB]B)
//#pragma oss task in([NB*NB]T) inout([NB*NB]B)
void dsyrk_tile( double *T, double *B, int NB)
{
    double const done=1.0, dmone=-1.0;
    cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, NB, NB, dmone, T, NB, done, B, NB);
}

void compute(int NB, int DIM, double ***A)
{
    int j,k,i;
    int N = NB*NB;
    for (j = 0; j < DIM; j++)
    {
        for (k= 0; k< j; k++)
        {
            for (i = j+1; i < DIM; i++)
            {
                dgemm_tile( A[k][i], A[k][j], A[j][i], NB);
            }
        }

        for (i = 0; i < j; i++)
        {
            dsyrk_tile( A[i][j], A[j][j], NB);
        }

        dpotrf_tile( A[j][j], NB);

        for (i = j+1; i < DIM; i++)
        {
            dtrsm_tile( A[j][j], A[j][i], NB);
        }
    }
    #pragma oss taskwait/*oss此处必须要有taskwait,否则check不通过*/
}

//--------------------------------------------------------------------------------

int check(int NB, int N, double *Alin, double ***A)
{
    int i,j;
    for (i = 0; i < N; i++)
    {
        for (j = 0; j < N; j++)
        {

            if(fabs(A[i/NB][j/NB][(i%NB)*NB+j%NB]-Alin[i*N+j])> 1e-4)
            {
                cout<<"check inconsistent!"<<endl;
                cout<<"in tile ["<<i/NB<<"]["<<j/NB<<"], element "<<(i%NB)*NB+j%NB<<" is wrong."<<endl;
                cout<<"the value in Tiled A is : "<<A[i/NB][j/NB][(i%NB)*NB+j%NB]<<endl;
                cout<<"but in original ALin, it is "<<Alin[i*N+j]<<endl;
                return 1;
            }
        }
    }
    return 0;
}

void randn(int n, double* Alin)
{
    srand((float)time(0));
    int i,j;
    double temp;
    for(i=0; i<n*n; i++)
    {
        Alin[i]=rand()%2;
    }
    for(i=0; i<n; i++)
    {
        for(j=0; j<n; j++)
        {
            temp=Alin[i*n+j]+Alin[j*n+j];
            Alin[i*n+j]=temp;
            Alin[j*n+i]=temp;
        }
    }
    for(i=0; i<n; i++)
    {
        Alin[i*n+i]=0;
    }
}

void init(int NB, int DIM)
{
    int N = NB*DIM;
    // linear matrix
    Alin = (double *) malloc(N * N * sizeof(double));
    // fill the matrix with random values
    randn(N, Alin);

    int i;
    int j;
    // make it positive definite
    for(i=0; i<N; i++)
    {
        Alin[i*N + i] += N;
    }
    // blocked matrix
    A = (double ***) malloc(DIM*sizeof(double **));
    for (i = 0; i < DIM; i++)
    {
        A[i] = (double **) malloc(DIM*sizeof(double*));
        for(j=0; j<DIM; j++)
            A[i][j]=(double *)malloc(NB*NB*sizeof(double));
    }

    for (i = 0; i < N; i++)
    {
        for (j = 0; j < N; j++)
        {
            A[i/NB][j/NB][(i%NB)*NB+j%NB] = Alin[i*N+j];
        }
    }
}

int main(int argc, char *argv[])
{
    if (argc<3)
    {
        cout<<"please input DIM NB\n"<<endl;
        exit(0);
    }

    int DIM = atoi(argv[1]);
    int NB = atoi(argv[2]);
    //int num_threads = atoi(argv[3]);

    struct timeval start;
    struct timeval stop;
    unsigned long elapsed;
    double exctime;

    init(NB, DIM);
    //omp_set_num_threads(num_threads);

    gettimeofday(&start,NULL);
   
        compute(NB, DIM, A);
    
    gettimeofday(&stop,NULL);

    elapsed = 1000000*(stop.tv_sec - start.tv_sec);
    elapsed += stop.tv_usec - start.tv_usec;
    exctime=(double)elapsed/1000000;
    cout<< "Operations time: "<<exctime<<endl;

    char LO='L';
    int  INFO;
    int N = NB*DIM;
    dpotrf_(&LO, &N, Alin, &N, &INFO);
    if(check(NB, N, Alin, A))
        exit(1);
    cout<<"check is passed, correct!"<<endl;

    return 0;
}
